package com.cloud.admin.support;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 并行任务执行辅助
 */
@Slf4j
public class ExecuteRunner {
    private static final ThreadLocal<CountDownLatch> countDownLatchHolder = new ThreadLocal<>();
    private static final ThreadPoolExecutor executor;

    // 初始化线程池
    static {
        int poolSize = Runtime.getRuntime().availableProcessors();
        executor = new ThreadPoolExecutor(poolSize, poolSize * 2, 180, TimeUnit.SECONDS, new LinkedBlockingQueue<>(poolSize), new ThreadPoolExecutor.CallerRunsPolicy());
    }

    /**
     * 执行开始，初始化上下文
     *
     * @param taskCount
     */
    public static void begin(int taskCount) {
        if (taskCount <= 0) {
            return;
        }
        countDownLatchHolder.set(new CountDownLatch(taskCount));
    }

    /**
     * 执行任务
     *
     * @param task
     * @return
     */
    public static void execute(Runnable task) {
        final CountDownLatch countDownLatch = countDownLatchHolder.get();
        executor.execute(() -> {
            try {
                task.run();
            } finally {
                if (countDownLatch != null) {
                    countDownLatch.countDown();
                }
            }
        });
    }

    /**
     * 等待所有任务都执行完毕
     */
    public static void await() {
        CountDownLatch countDownLatch = countDownLatchHolder.get();
        if (countDownLatch != null) {
            try {
                countDownLatch.await();
            } catch (InterruptedException e) {
                log.error(e.getMessage(), e);
            }
            countDownLatchHolder.remove();
        }
    }
}
